#!/usr/bin/python

import json
import os.path
import math

from pprint import PrettyPrinter
from sys import argv

pp = PrettyPrinter(indent=4)

'''
A couple of observations:
    rocket packs in HG ignore the acc modifier when calculating their ammo cost
    autocannons, rocket packs use 2*ROF for the rofCost for some reason
    guided weapons use guided-ammo instead of the weapon system proper having guided
    
'''

# Accuracy Multiplier 
aDict = {
    5 : 180.0,  4 : 36.0,  3 : 9.0,  2 : 3.0,    1 : 1.5, 
    0 : 1.0, 
   -1 : 0.667, -2 : 0.5,  -3 : 0.4, -4 : 0.333, -5 : 0.286
}
    
debug = True
jccCalcs = False

# TODO: Weapons ranges at 0

class Characteristic:
    def __init__(self, name, dict):
        self.name = name
        self.desc = ""
        self.multi = 0.0
        self.a_multi = 0.0
        self.g_multi = 0.0
        self.w_multi = 0.0

        for key,value in dict.iteritems():
            if key == 'multiplier':
                self.multi = value
            elif key == 'ammoCostMulti':
                self.a_multi = value
            elif key == 'generalCostMulti':
                self.g_multi = value
            elif key == 'weaponCostMulti':
                self.w_multi = value
            elif key == 'desc':
                self.desc = value

    def getWeaponCost(self, weapon, baseCost):
        multi = self.w_multi
        if self.g_multi > 0:
            multi = self.g_multi

        if self.name == 'anti-infantry':
            multi = 1.0 + (0.1  * weapon.dam)
        elif self.name.find('minimum-range') != -1:
            multi = self.w_multi / weapon.range_base
            #print "w_multi (%2.2f) / range_base (%2.2f) = (%2.2f)" % (self.w_multi, weapon.range_base, multi)
            multi = 1.0 - (0.1 * multi)
        elif self.name.find('attenuating-damage') != -1:
            multi = 1.0

        #print "        %s multi is: %2.2f" % (self.name, multi)
        return baseCost * multi

    def getAmmoCost(self, weapon, baseCost):
        multi = self.g_multi
        if self.a_multi != 0:
            multi = self.a_multi

        return baseCost * multi

class Weapon:

    def __init__(self, dict):
        self.label = ""
        self.desc = ""

        self.dam = 0.0
        self.calcDam = 0.0 # Used for attenuating damage weapons

        self.range_base = 0
        self.range_med = 0
        self.range_long = 0
        self.range_ex = 0
    
        self.acc = 0
        self.acc_multi = 1
        self.rof = 0
    
        self.characteristics = []

        self.isSingle = True
        self.isMultiFunc = False
        self.isMated = False
        self.isSplit = False
        self.refLabels = []

        self.baseCost = 0
        self.baseAmmoCost = 0
        self.modCost = 0
        self.modAmmoCost = 0
        self.finalTV = 0
        self.finalAmmoTV = 0
        self.oldCost = 0
        self.oldAmmoCost = 0

        self.baseMinSize = 1
        self.finalMinSize = 1
    
        self.isHEAT = False
        self.isMelee = False
        self.isThrown = False
        self.isInfantry = False
        self.hasAmmo = True

        for key,value in dict.iteritems():
            if key == 'label':
                self.label = value
            elif key == 'type':
                if value == 'single':
                    continue
                self.isSingle = False
                self.refLables = dict['referenceLabels']

                if value == 'multiFunction':
                    isMultiFunc = True
                elif value == 'mated':
                    isMultiFunc = True
                elif value == 'split':
                    isSplit = True
            elif key == 'description':
                self.desc = value
                #print "Desc is: ", self.desc
            elif key == 'damageMulti':
                self.dam = value
                if self.calcDam != 0:
                    self.calcDam= value + self.calcDam
                else:
                    self.calcDam = self.dam
                #print "DAM is x%i" % (self.dam)
            elif key == 'baseRange':
                self.range_base = value
                self.range_med = value * 2
                self.range_long = value * 4
                self.range_ex = value * 8 
                #print "Weapon ranges are: %i/%i/%i/%i" % (self.range_base, self.range_med, self.range_long, self.range_ex)
            elif key == 'accuracy':
                self.acc = value
                self.acc_multi = aDict[value]
                #print "ACC Multi is: %2.2f for Acc: %i" % (self.acc_multi, self.acc)
            elif key == 'rateOfFire':
                self.rof = value
                #print "RoF is: %f" % (self.rof)
            elif key == 'finalMinimumSize':
                self.finalMinSize = value
                #print "Final Min Size will be: %i" % (self.finalMinSize)
            elif key == 'isHEAT':
                self.isHEAT = value
            elif key == 'isMelee':
                self.isMelee = value
            elif key == 'isThrown':
                self.isThrown = value
            elif key == 'isInfantry':
                self.isInfantry = value
            elif key == 'oldCost':
                self.oldCost = value
            elif key == 'oldAmmoCost':
                self.oldAmmoCost = value
            elif key == 'characteristics':
                self.characteristics = value
                for key in value:
                    #print "Key is: ",key
                    if key == 'self-destruct':
                        self.hasAmmo = False
                    elif key.find('attenuating-damage-') != -1:
                        mod = key[key.rfind('-')+1:]
                        #print "Attenuating-damage value: ", mod
                        self.calcDam = self.calcDam - float(mod) 
                        #print "self.calcDam (%2.2f) vs. (%2.2f)" % (self.calcDam, self.dam)

        #print "Done loading."

    def __cmp__(self, dict):
        if not isinstance(dict, Weapon):
            return cmp(self.name, dict)

        if self.isMelee and not dict.isMelee:
            return -1
        if not self.isMelee and dict.isMelee:
            return 1
        
        return 0
            

def loadCharacteristics(path):
    # Try to load from the subdirectory first
    charPath = os.path.join(path, "characteristics.json")
    if os.path.exists(charPath):
      f = open(charPath, 'r')
    else:
      f = open("characteristics.json", 'r')

    tmpDict = json.load(f)

    cDict = {}
    for key, values in tmpDict.iteritems():
        x = Characteristic(key, values)
        cDict[x.name] = x 
        
    return cDict

def calculateWeapon(wDict, cDict, debug=False):
    # Do asserts here for the values we need?

    weapon = Weapon(wDict)
    
    rofCost = 0.0
    if weapon.rof > 0:
        if jccCalcs:
          # JCC Formula
          rofCost = (0.5 * (weapon.calcDam + (weapon.rof)))**2
        else:
          # HG Formula
          rofCost = (0.5 * (weapon.calcDam + (weapon.rof*2)))**2
        if debug:
            print "  rofCost = ((0.5 * (%i + %i))**2 = %2.2f" % (weapon.calcDam, weapon.rof, rofCost)
            print "rofCost      is: %5.2f" % (rofCost)
    
    # Calc base costs
    baseCost = 0.0
    baseAmmoCost = 0.0
    if debug:
        print ""
    if weapon.isMelee:
        baseCost     = weapon.acc_multi * ( (weapon.calcDam - 2.0)**2 + rofCost)
        if debug:
            print "  melee weapon cost is: (%2.2f * ((%i - 2.0)**2 + %2.2f)) = %2.2f" % \
                (weapon.acc_multi, weapon.calcDam, rofCost, baseCost)
    else:
        baseCost     = weapon.acc_multi * ( weapon.calcDam**2 + weapon.range_base**3 + rofCost)
        if debug:
            print "  ranged weapon cost is: (%2.2f * ((%i)**2 + (%i)**3 + %2.2f) = %2.2f" % \
                (weapon.acc_multi, weapon.calcDam, weapon.range_base, rofCost, baseCost)
    modCost = getWeaponModifiedCost(baseCost, weapon, cDict, debug)
    modCost = math.ceil(modCost)
    if debug:
        print "baseCost     is: %5.2f" % (baseCost)
        print "modCost      is: %5.2f" % (modCost)
        print ""

    if (weapon.hasAmmo):
        if weapon.isMelee:
            if jccCalcs:
              # JCC
              baseAmmoCost = weapon.acc_multi * ( (weapon.calcDam - 2.0)**2)
            else:
              # HG Seems to drop the weapon.acc_multi calc here...
              baseAmmoCost = (weapon.calcDam - 2.0)**2
            if debug:
                print "  base ammo cost is: (%2.2f * ((%i - 2.0)**2) = %2.2f / 250.0" % \
                    (weapon.acc_multi, weapon.calcDam, baseAmmoCost)
        else:
            if jccCalcs:
              # JCC
              baseAmmoCost = weapon.acc_multi * ( weapon.calcDam**2 + weapon.range_base**3)
            else:
              # HG seems to drop the weapon.acc_multi calc here
              baseAmmoCost = ( weapon.calcDam**2 + weapon.range_base**3)
            if debug:
                print "  base ammo cost is: (%2.2f * ((%i)**2 + (%i)**3 ) = %2.2f / 250.0" % \
                    (weapon.acc_multi, weapon.calcDam, weapon.range_base, baseAmmoCost)
        baseAmmoCost = baseAmmoCost / 250.0
        modAmmoCost = getAmmoModifiedCost(baseAmmoCost, weapon, cDict, debug)
        if debug:
            print "baseAmmoCost is: %5.2f" % (baseAmmoCost)
            print "modAmmoCost  is: %5.2f" % (modAmmoCost)

    # Minimum size
    baseMinSize = 0.0
    if weapon.isHEAT:
        baseMinSize = ((modCost / 2.0)**(1.0/3.0)) * 0.75
        if debug:
            print ""
            print "  Weapon is HEAT, reducing baseMinSize to %2.2f" % (baseMinSize)
    else:
        baseMinSize = (modCost / 2.0)**(1.0/3.0)
    #baseMinSize = math.floor(baseMinSize)
    if debug:
        print "Raw baseMinSize = (%2.2f)" % (baseMinSize)
    baseMinSize = round(baseMinSize)
    if debug:
        print "baseMinSize  is: %5.2f" % (baseMinSize)

    if weapon.finalMinSize != 0 and weapon.finalMinSize < baseMinSize:
        threatValue = (modCost * baseMinSize) / weapon.finalMinSize
        if debug:
            print "  Miniaturizing weapon from (%i @ %2.2f) to (%i @ %2.2f)" % (baseMinSize, modCost, weapon.finalMinSize, threatValue)
    else:
        threatValue = modCost

    # Round values to two significant digits
    weapon.finalTV = round(threatValue, 2)
    if debug:
        print ""
        print "finalTV      is: %2.2f" % (weapon.finalTV)

    if weapon.oldCost != 0:
        deltaTV = weapon.finalTV - weapon.oldCost
        deltaP = (deltaTV / weapon.oldCost) * 100.0
        if debug:
            print "Sanity Check: Weapon TV [%i TV versus %i TV] = [%+2.2f TV / %+2.2f%%] delta" %  \
                (weapon.finalTV, weapon.oldCost, deltaTV, deltaP) 

    if (weapon.hasAmmo):
        weapon.finalAmmoTV    = round(modAmmoCost, 2)
        if debug:
            print ""
            print "ammoTV       is: %2.2f" % (weapon.finalAmmoTV)
        if weapon.oldAmmoCost != 0:
            deltaTV = weapon.finalAmmoTV - weapon.oldAmmoCost
            deltaP = (deltaTV / weapon.oldAmmoCost) * 100.0
            if debug:
                print "Sanity Check: Ammo TV [%2.2f TV versus %2.2f TV] = [%+2.2f TV / %+2.2f%%] delta" %  \
                    (weapon.finalAmmoTV, weapon.oldAmmoCost, deltaTV, deltaP)
    return weapon

def getWeaponModifiedCost(baseCost, weapon, cDict, debug=False):
    workingCost = baseCost
    for cname in weapon.characteristics:
        character = cDict[cname]
        tmpCost = character.getWeaponCost(weapon, workingCost)
        if tmpCost != 0 :
            if debug:
                multi = tmpCost / workingCost  
                print "  Modify workingCost to (%2.2f) due to [%s] (x%s)" % (tmpCost, character.name, multi)
            workingCost = tmpCost
        
    return workingCost

def getAmmoModifiedCost(baseCost, weapon, cDict, debug=False):
    workingCost = baseCost
    for cname in weapon.characteristics:
        character = cDict[cname]
        tmpCost = character.getAmmoCost(weapon, workingCost)
        if tmpCost != 0:
            if debug:
                multi = tmpCost / workingCost
                print "  Modify ammo workingCost to (%2.2f) due to [%s] (x%s)" % (tmpCost, character.name, multi)
            workingCost = tmpCost
    return workingCost

def readWeapon(arg, dirname, names):
    for name in names:
        if name.endswith('.swp') or name.endswith('.sortorder'):
          continue

        fullPath = os.path.join(dirname, name)
        if os.path.isdir(fullPath):
            continue
        if debug:
            print "Opening file: [%s]" % (fullPath)
        f = open(fullPath, 'r')
        wDict = json.load(f)
        weapon = calculateWeapon(wDict, cDict, debug)
        weaponsDict[weapon.label] = weapon

if __name__ == '__main__':
    cDict = loadCharacteristics(argv[1]);
    if os.path.isdir(argv[1]):
        # Load the sorting order
        so = open("%s/%s.sortorder" % (argv[1], argv[1]), 'r')
        sortL = json.load(so)

        weaponsDict = {}
        # Batch mode - process all files, output in csv format
        os.path.walk(argv[1], readWeapon, (cDict, weaponsDict, debug))

        sortList = [ sortL['cannons'], sortL['missiles'], sortL['recoilless'], \
                    sortL['support'], sortL['advanced'], \
                    sortL['thrown'], sortL['melee'] ]
        if 'infantry' in sortL:
          sortList.append(sortL['infantry'])
        
        # Human readable
        if debug:
            for category in sortList:
                for label in category:
                    if label.upper() not in weaponsDict:
                      continue
                    weapon = weaponsDict[label.upper()]
                    print "  %8s  %-32s  %8i  %2i / %2i / %2i / %2i  x%-2i  %+2i %2i %6.2f   %2i  %s" % \
                        (weapon.label, weapon.desc, weapon.finalTV, \
                        weapon.range_base, weapon.range_med, weapon.range_long, weapon.range_ex, \
                        weapon.dam, weapon.acc, weapon.rof, weapon.finalAmmoTV, weapon.finalMinSize, \
                        ",".join(weapon.characteristics))

        # HGB-TV-CALC Weapons format
        w = open("%s.weapons.xml" % argv[1], 'w')
        w.write('''<?xml version="1" encoding="ISO-8859-1" ?><weaponData>\n''')

        for category in sortList:
            for label in category:
                if label.upper() not in weaponsDict:
                  continue
                weapon = weaponsDict[label.upper()]
                weaponL = []
                weaponL.append( "<weapon id=\"%s\" rating=\"%s\" ammo=\"%s\" msize=\"%s\" dam=\"%s\" brange=\"%s\"" % \
                    (weapon.label.lower(), weapon.finalTV, weapon.finalAmmoTV, weapon.finalMinSize, weapon.dam, weapon.range_ex)
                    )

                if weapon.acc > 1:
                    weaponL.append("accurate=\"1\"")
                if weapon.acc < 0:
                    weaponL.append("inaccurate=\"1\"")
                if weapon.isHEAT:
                    weaponL.append("heat=\"1\"")
                if weapon.isMelee:
                    weaponL.append("melee=\"1\"")
                if weapon.isThrown:
                    weaponL.append("thrown=\"1\"")
                if weapon.isInfantry: 
                    weaponL.append("infantry=\"1\"")
                if weapon.label.find("AA") != -1:
                    weaponL.append("aa=\"1\"")
                if weapon.rof > 0:
                    weaponL.append("rof=\"%s\"" % weapon.rof)
                    if 'missile' in weapon.characteristics:
                        weaponL.append(" barrage=\"%s\"" % weapon.rof)
                    else:
                        weaponL.append(" burst=\"%s\"" % weapon.rof)

                for c in weapon.characteristics:
                    if c == 'anti-infantry' or c == 'anti-infantry-ammo':
                        weaponL.append( "ai=\"1\"" )
                    elif c == 'area-effect-0h-5m' or c == 'area-effect-0h-5m-ammo':
                        weaponL.append( "bae=\"1\"" )
                    elif c == 'area-effect-0h-10m' or c == 'area-effect-0h-10m-ammo':
                        weaponL.append( "bae=\"2\"" )
                    elif c == 'area-effect-0h-15m' or c == 'area-effect-0h-15m-ammo':
                        weaponL.append( "bae=\"3\"" )
                    elif c == 'area-effect-0h-20m' or c == 'area-effect-0h-20m-ammo':
                        weaponL.append( "bae=\"4\"" )
                    elif c == 'armor-piercing':
                        weaponL.append( "ap=\"1\"" );
                    elif c == 'armor-crushing':
                        weaponL.append( "ac=\"1\"" );
                    elif c == 'attenutating-damage-1':
                        weaponL.append( "mb=\"1\"" )
                    elif c == 'attenutating-damage-2':
                        weaponL.append( "mb=\"2\"" )
                    elif c == 'attenutating-damage-3':
                        weaponL.append( "mb=\"3\"" )
                    elif c == 'attenutating-damage-4':
                        weaponL.append( "mb=\"4\"" )
                    elif c == 'guided' or c == 'guided-ammo':
                        weaponL.append( "guided=\"1\"" )
                    elif c == 'haywire' or c == 'haywire-ammo':
                        weaponL.append( "haywire=\"1\"" );
                    elif c == 'incendiary' or c == 'incendiary ammo':
                        weaponL.append( "sb=\"1\"" )
                    elif c == 'indirect-fire':
                        weaponL.append( "if=\"1\"" )
                    elif c.find('minimum-range') != -1:
                        weaponL.append( "mr=\"1\"" )
                    elif c.find('scatter') != -1:
                        weaponL.append( "scatter=\"1\"" )
                    elif c.find('attenuating-damage') != -1:
                        weaponL.append( "mb=\"%s\"" % c[c.rfind("-")+1:] )
                    
                w.write("%s />\n" % (" ".join(weaponL)))
        w.write('''</weaponData>\n''')
        w.close()

    else:
        readWeapon((cDict, debug), 
            os.path.dirname(argv[1]), 
            [os.path.basename(argv[1])]
            )
